"""
群体数据与亲本进行比对，并生成格式数据
"""

import sys
import pandas as pd

vcf_file_path= sys.argv[1]
parent=sys.argv[2]
out_dir=sys.argv[3]


file_handle = open(vcf_file_path, "r")

line_num = 0
while True:
    line = file_handle.readline()
    if "#CHROM" in line:
        break
    else:
        line_num+=1

# 打开文件
vcf_df = pd.read_csv(vcf_file_path, skiprows=line_num, sep="\t", low_memory=False)

samples = vcf_df.columns.tolist()[9:]
parent_data = vcf_df[parent].tolist()


def format_data(data):
    new_data = []
    for index, value in enumerate(data):
        if value[0] != value[-1]:
            new_data.append('h')
        elif value == parent_data[index]:
            new_data.append('a')
        else:
            new_data.append('b')
    return new_data


for sample in samples:
    sample_data = vcf_df[sample].tolist()
    vcf_df[sample] = format_data(sample_data)

groups = vcf_df.groupby("#CHROM")

for chr_id, chr_df in groups:
    new_chr_df = chr_df[["#CHROM", "POS"]+samples].copy()
    new_chr_df["#CHROM"] = new_chr_df.apply(lambda x: f"{x['#CHROM']}_{x['POS']}", axis=1)

    new_chr_df.rename(
        columns={"#CHROM":"marker", "POS":"position(bp)"},
        inplace=True
    )
    new_chr_df.to_csv(f"{out_dir}/crosspoints_{chr_id}.tsv", index=False, sep="\t")
